import math
import numpy as np
from copy import copy

from History import *
from ValidityInterval import *

modelID = 0

def interpolateModels(m1, m2, state, nextState, dt):
    newModel = Model2()
    # 
    return newModel
    


class Model2:
    def __init__(self):
        global modelID
        self.k1 = 0
        self.k2 = 0
        self.k3 = 0
        self.validityInterval = ValidityInterval(4) # Number of dimension in a state
        self.goodStates = []
        self.badStates = []
        self.id = modelID
        modelID = modelID + 1
        
    def limitControls(self, c1, c2):
        if c1 < -1:
            c1 = -1
        elif c1 > 10:
            c1 = 10
        if c2 < -1:
            c2 = -1
        elif c2 > 10:
            c2 = 10
        return (c1, c2)
     
    def getControls(self, dt, a, aDot, aNew, aDotNew, stiffness):
        c1 = (aDotNew - aDot - a*dt*self.k3 - 2.0*stiffness*dt*self.k2/3.0)
        c1 = c1 / (dt*self.k1 + dt*self.k2/3.0)
        c2 = (2.0*stiffness + c1) / 3.0
        # Test that condition is met
        if 0.5*(c1+c2) - math.fabs(c1-c2) - stiffness < 1e-8:
            return self.limitControls(c1, c2)
        c1 = (aDotNew - aDot - a*dt*self.k3 + 2.0*stiffness*dt*self.k2)
        c1 = c1 / (dt*self.k1 + 3.0*dt*self.k2)
        c2 = 3.0*c1 - 2.0*stiffness
        if 0.5*(c1+c2) - math.fabs(c1-c2) - stiffness < 1e-8:
            return self.limitControls(c1, c2)
        return (0,0)
    
    def setParameters(self, events):
        aRows = []
        RHSs = []
        ep = events[0]
        for e in events[1:]:
            dt = e.time - ep.time
            RHSs.append(e.state.arm1.aDot - ep.state.arm1.aDot)
            a1 = ep.state.arm1.c1 * dt
            a2 = ep.state.arm1.c2 * dt
            a3 = ep.state.arm1.a * dt
            aRows.append([a1, a2, a3])
            ep = e
        #print RHSs
        #print aRows
        rhs = np.array(RHSs)
        a = np.array(aRows)
        x = np.linalg.lstsq(a, rhs)
        self.k1 = x[0][0]
        self.k2 = x[0][1]
        self.k3 = x[0][2]
        print "Model: Set parameters {0}".format(self)

    def setGoodStates(self, events):
        ep = events[0]
        for e in events[1:]:
            # Is ep already in good states?
            found = False
            for state in self.goodStates:
                if stateDistance(ep.state, state) < 0.01:
                    #print "Model {0}: Good state already present: {1}".format(self.id, ep.state)
                    found = True
                    break
            if found:
                ep = e
                continue
            a = ep.state.arm1.a
            aDot = ep.state.arm1.aDot
            dt = e.time - ep.time
            aDotNew = self.k1*ep.state.arm1.c1*dt + self.k2*ep.state.arm1.c2*dt
            aDotNew = aDotNew + self.k3*ep.state.arm1.a*dt + ep.state.arm1.aDot
            aNew = 0.5*(aDot+aDotNew)*dt + a
            # print distance
            #print "a   : " + str(aNew) + "   " + str(e.state.arm1angle) + "   " + str(aNew - e.state.arm1angle)
            #print "aDot: " + str(aDotNew) + "   " + str(e.state.arm1angularVelocity) + "   " + str(aDotNew - e.state.arm1angularVelocity)
            if math.fabs(aNew - e.state.arm1.a)/dt < 1 and math.fabs(aDotNew - e.state.arm1.aDot)/dt < 1:
                print "Model {0}: Found good state: {1}".format(self.id, ep.state)
                ep.state.modelExists = True
                self.addGoodState(ep.state)
            ep = e
    
    def minDistanceToGoodAndBadState(self, state, nextState, dt, stiffness):
        dGood = 1000000
        dBad = 1000000
        s = copy(state)
        (s.arm1.c1, s.arm1.c2) = \
               self.getControls(dt, s.arm1.a, s.arm1.aDot, nextState.arm1.a, nextState.arm1.aDot, stiffness)
        for goodState in self.goodStates:
            print "Model {0}: Good: {1}".format(self.id, goodState)
            dN = stateDistance(s, goodState)
            if dN < dGood:
                dGood = dN
        for badState in self.badStates:
            print "Model {0}: Bad: {1}".format(self.id, badState)
            dN = stateDistance(s, badState)
            if dN < dBad:
                dBad = dN
        return (dGood, dBad)
    
    def addBadState(self, state):
        self.badStates.append(copy(state))
        
    def addGoodState(self, state):
        self.goodStates.append(copy(state))

    def __str__(self):
        return "k1: {0: 6.2f}, k2: {1: 6.2f}, k3: {2: 6.2f}".format(self.k1, self.k2, self.k3)

    def printMe(self):
        print "Model: k1 = " + str(self.k1)
        print "Model: k2 = " + str(self.k2)
        print "Model: k3 = " + str(self.k3)
        print "Model {0}: Good states:".format(self.id)
        for s in self.goodStates:
            print "Model:   {0}".format(s)
        print "Model {0}: Bad states:".format(self.id)
        for s in self.badStates:
            print "Model:   {0}".format(s)
